from collections import Counter
from math import log

hmm_model = {i: Counter() for i in 'sbme'}

with open('dict.txt') as f:
    for line in f:
        lines = line.split(' ')
        if len(lines[0]) == 1:
            hmm_model['s'][lines[0]] += int(lines[1])
        else:
            hmm_model['b'][lines[0][0]] += int(lines[1])
            hmm_model['e'][lines[0][-1]] += int(lines[1])
        for m in lines[0][1:-1]:
            hmm_model['m'][m] += int(lines[1])

log_total = {i: log(sum(hmm_model[i].values())) for i in 'sbme'}

trans = {'ss': 0.3,
         'sb': 0.7,
         'bm': 0.3,
         'be': 0.7,
         'mm': 0.3,
         'me': 0.7,
         'es': 0.3,
         'eb': 0.7
         }

trans = {i: log(j) for i, j in trans.items()}


def viterbi(nodes):
    paths = nodes[0]
    for l in range(1, len(nodes)):
        paths_ = paths
        paths = {}
        for i in nodes[l]:
            nows = {}
            for j in paths_:
                if j[-1]+i in trans:
                    nows[j+i] = paths_[j]+nodes[l][i]+trans[j[-1]+i]

            nows_values_list = list(nows.values())
            k = nows_values_list.index(max(nows_values_list))

            nows_keys_list = list(nows.keys())
            paths[nows_keys_list[k]] = nows_values_list[k]

    paths_values_list = list(paths.values())
    paths_keys_list = list(paths.keys())

    return paths_keys_list[paths_values_list.index(max(paths_values_list))]


def hmm_cut(s):
    nodes = [{i: log(j[t]+1)-log_total[i]
                for i, j in hmm_model.items()} for t in s]
    tags = viterbi(nodes)
    words = [s[0]]
    for i in range(1, len(s)):
        if tags[i] in ['b', 's']:
            words.append(s[i])
        else:
            words[-1] += s[i]
    return words
